"""
创建 SIMION .fly2 粒子定义文件
圆形分布，使用Lua自定义分布函数
参考: https://simion.com/info/fly2_file.html
"""

import os


def create_fly2_circle(
    n=5000,
    mass=28,
    charge=1,
    cx=0,
    cy=0,
    z=0,
    diameter=10,
    ke=90000,
    fill=True,          # True=填充整个圆, False=只在圆周边缘
    output_file=None
):
    """
    创建圆形分布的粒子文件（使用Lua自定义分布）
    
    参考: https://simion.com/info/fly2_file.html
    使用 distribution() 函数和自定义Lua函数实现圆形分布
    """
    
    radius = diameter / 2.0
    
    lines = []
    lines.append("-- SIMION FLY2 粒子定义文件")
    lines.append(f"-- 圆形分布, 直径={diameter}mm, 中心=({cx},{cy},{z})")
    lines.append("-- 使用Lua自定义分布函数")
    lines.append("")
    lines.append("-- 圆形分布参数")
    lines.append(f"local cx = {cx}      -- 圆心x")
    lines.append(f"local cy = {cy}      -- 圆心y")
    lines.append(f"local radius = {radius}  -- 半径")
    lines.append("")
    
    # 使用共享变量确保x和y来自同一个随机点
    lines.append("-- 共享变量存储当前点的坐标")
    lines.append("local current_x, current_y = cx, cy")
    lines.append("local need_new_point = true")
    lines.append("")
    
    if fill:
        # 填充整个圆的均匀分布
        lines.append("-- 生成圆内均匀分布的新点")
        lines.append("function generate_point()")
        lines.append("  local r = radius * math.sqrt(simion.rand())")
        lines.append("  local theta = simion.rand() * 2 * math.pi")
        lines.append("  current_x = cx + r * math.cos(theta)")
        lines.append("  current_y = cy + r * math.sin(theta)")
        lines.append("end")
    else:
        # 仅圆周边缘
        lines.append("-- 生成圆周边缘的新点")
        lines.append("function generate_point()")
        lines.append("  local theta = simion.rand() * 2 * math.pi")
        lines.append("  current_x = cx + radius * math.cos(theta)")
        lines.append("  current_y = cy + radius * math.sin(theta)")
        lines.append("end")
    
    lines.append("")
    lines.append("-- x坐标分布函数 (生成新点)")
    lines.append("function circle_x()")
    lines.append("  generate_point()")
    lines.append("  return current_x")
    lines.append("end")
    lines.append("")
    lines.append("-- y坐标分布函数 (使用已生成的点)")
    lines.append("function circle_y()")
    lines.append("  return current_y")
    lines.append("end")
    
    lines.append("")
    lines.append("particles {")
    lines.append("  coordinates = 0,")
    lines.append("  standard_beam {")
    lines.append(f"    n = {n},")
    lines.append("    tob = arithmetic_sequence {")
    lines.append("      first = 0,")
    lines.append("      step = 1")
    lines.append("    },")
    lines.append(f"    mass = {mass},")
    lines.append(f"    charge = {charge},")
    lines.append("    x = distribution(circle_x),")
    lines.append("    y = distribution(circle_y),")
    lines.append(f"    z = {z},")
    lines.append(f"    ke = {ke},")
    lines.append("    cwf = 1,")
    lines.append("    color = 0,")
    lines.append("    direction = vector(0, 0, 1)")
    lines.append("  },")
    lines.append("}")
    
    fly2_content = "\n".join(lines)
    
    # 保存文件
    script_dir = os.path.dirname(os.path.abspath(__file__))
    if output_file is None:
        fill_tag = "fill" if fill else "ring"
        output_file = os.path.join(
            script_dir,
            f'particles_n{n}_m{mass}_q{charge}_circle_d{diameter}mm_{fill_tag}.fly2'
        )
    elif not os.path.isabs(output_file):
        output_file = os.path.join(script_dir, output_file)
    
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(fly2_content)
    
    print(f"粒子文件已保存: {output_file}")
    print(f"  粒子数量: {n}")
    print(f"  质量: {mass} amu")
    print(f"  电荷: {charge}")
    print(f"  动能: {ke} eV")
    print(f"  圆形分布直径: {diameter} mm")
    print(f"  填充模式: {'填充整个圆' if fill else '仅圆周边缘'}")
    print(f"  中心: ({cx}, {cy}, {z})")
    
    return output_file


if __name__ == '__main__':
    import argparse
    
    parser = argparse.ArgumentParser(description='创建SIMION .fly2粒子文件 (圆形分布)')
    parser.add_argument('-n', type=int, default=5000, help='粒子数量')
    parser.add_argument('-m', '--mass', type=float, default=28, help='质量 (amu)')
    parser.add_argument('-q', '--charge', type=int, default=1, help='电荷')
    parser.add_argument('-x', type=float, default=0, help='中心x坐标 (mm)')
    parser.add_argument('-y', type=float, default=0, help='中心y坐标 (mm)')
    parser.add_argument('-z', type=float, default=0, help='z坐标 (mm)')
    parser.add_argument('-d', '--diameter', type=float, default=10, help='圆形直径 (mm)')
    parser.add_argument('--ke', type=float, default=90000, help='动能 (eV)')
    parser.add_argument('--ring', action='store_true', help='仅圆周边缘分布 (默认填充整个圆)')
    parser.add_argument('-o', '--output', type=str, default=None, help='输出文件名')
    
    args = parser.parse_args()
    
    create_fly2_circle(
        n=args.n,
        mass=args.mass,
        charge=args.charge,
        cx=args.x,
        cy=args.y,
        z=args.z,
        diameter=args.diameter,
        ke=args.ke,
        fill=not args.ring,
        output_file=args.output
    )
    # python create_fly2.py -n 5000 -m 28 -q 1 -x 61 -y 61 -d 10
